#include "IR_optimizer.h"
#include "Function.h"
#include "GlobalVariable.h"
#include "ConstInt.h"
#include "IRCode.h"
#include "Instruction.h"
#include <algorithm>
#include <unordered_set>
#include <unordered_map>
#include "Module.h"
#include "LocalVariable.h"
class Module;
class Function;
class Instruction;
class GlobalVariable;
class LocalVariable;
IR_optimizer::IR_optimizer(Module *_module) : module(_module) {}

static Value *computeConstExpr(IRInstOperator op, Value *op1, Value *op2);

void IR_optimizer::optimizeAll() {
    // 执行死代码消除
    deadCodeElimination();
    constantPropagation();
    deadCodeElimination();
    printf("执行优化\n");
    // 其他优化暂时不实现
    // optimizeGlobalVariables();
    // optimizeFunctions();
    // functionInlining();
    // registerAllocationOptimization();
}
void IR_optimizer::constantPropagation() {
    for (auto *func : module->getFunctionList()) {
        if (func->isBuiltin()) continue;

        auto &interCode = func->getInterCode();
        std::unordered_map<Value*, ConstInt*> constMap;
        std::vector<Instruction*> toRemove;
        std::unordered_map<Value*, Value*> replacedWith; 
        int iteration = 1;
        bool changed;

        do {
            changed = false;
            toRemove.clear();

            printf("第%d次迭代\n", iteration++);

            for (auto *inst : interCode.getInsts()) {
                if (inst->getOp() == IRInstOperator::IRINST_OP_ASSIGN) {
                    Value *dst = inst->getOperand(0);
                    Value *src = inst->getOperand(1);

                    if (auto it = constMap.find(src); it != constMap.end()) {
                        src = it->second;
                    }
                    
                    if (!dynamic_cast<ConstInt*>(src) && dynamic_cast<ConstInt*>(dst)) {
                        constMap.erase(replacedWith[dst]);
                        if (replacedWith.count(dst)) {
                            printf("进入\n");
                            Value *oldval = replacedWith[dst];
                            inst->setOperand(0, oldval);  // 恢复
                            printf("警告:call 结果写入常量变量 %s，恢复原变量\n", replacedWith[dst]->getName().c_str());
                            printf("inst %s 的操作数 %d 恢复为 %s\n", inst->getName().c_str(), 0, inst->getOperand(0)->getName().c_str());
                            }
                        changed = true;  // 标记为已更改
                        break;
                    }

                    if (auto *cval = dynamic_cast<ConstInt*>(src)) {
                        // 清除旧值（防止多次赋值）
                        constMap.erase(dst);
                        constMap[dst] = cval;
                        bool isReplaced = replacedWith.count(dst);
                        Value *oldVal = nullptr;
                        if(isReplaced){
                            oldVal = replacedWith[dst];
                            replacedWith.erase(dst);
                            replacedWith[cval] = oldVal;  // 更新映射
                        }else{
                            replacedWith[cval] = dst;  // 记录替换关系
                            printf("replacedWith[%s] = %s\n", cval->getName().c_str(), dst->getName().c_str());
                        }
                        toRemove.push_back(inst);
                        changed = true;
                        printf("常量折叠: %s = %d\n", isReplaced ? oldVal->getName().c_str() : dst->getName().c_str(), cval->getVal());
                    }

                } else if (inst->isBinaryOp()) {
                    Value *dst = inst->getOperand(2);
                    Value *op1 = inst->getOperand(0);
                    Value *op2 = inst->getOperand(1);

                    Value* resolvedOp1 = constMap.count(op1) ? constMap[op1] : op1;
                    Value* resolvedOp2 = constMap.count(op2) ? constMap[op2] : op2;

                    if (auto *c1 = dynamic_cast<ConstInt*>(resolvedOp1)) {
                        if (auto *c2 = dynamic_cast<ConstInt*>(resolvedOp2)) {
                            Value *res = computeConstExpr(inst->getOp(), c1, c2);
                            if (auto *cres = dynamic_cast<ConstInt*>(res)) {
                                constMap.erase(dst);
                                constMap[dst] = cres;
                                bool isReplaced = replacedWith.count(dst);
                                Value *oldVal = nullptr;
                                if(isReplaced){
                                    // 如果 dst 已经被替换过，清除旧的映射
                                    oldVal = replacedWith[dst];
                                    replacedWith.erase(dst);
                                    replacedWith[cres] = oldVal;  // 更新映射
                                }else{
                                    replacedWith[cres] = dst;  // 记录替换关系
                                    printf("replacedWith[%s] = %s\n", cres->getName().c_str(), dst->getName().c_str());
                                }
                                toRemove.push_back(inst);
                                changed = true;
                                printf("常量折叠: %s = %d\n", isReplaced ? oldVal->getName().c_str() : dst->getName().c_str(), cres->getVal());
                            }
                        }
                    }
                } 
            }

            // 使用点替换（替换后不会立即生效，因为本轮中还要继续迭代）
            for (auto &[val, cval] : constMap) {
                for (auto *use : val->getUseList()) {
                    // 仅替换不在定义指令前的 use
                    if (auto *inst = dynamic_cast<Instruction*>(use->getUser())) {
                        // 替换操作数 val 为 cval
                        for (int i = 0; i < inst->getOperandsNum(); ++i) {
                            if (inst->getOperand(i) == val) {
                                inst->setOperand(i, cval);
                            }
                        }
                    }
                }
            }

            // 删除已处理的指令
            auto &insts = interCode.getInsts();
            insts.erase(std::remove_if(insts.begin(), insts.end(),
                [&](Instruction *i) {
                    return std::find(toRemove.begin(), toRemove.end(), i) != toRemove.end();
                }), insts.end());

        } while (changed);
    }
}


// 辅助函数：计算两个常量的运算结果
Value *computeConstExpr(IRInstOperator op, Value *op1, Value *op2) {
    if (!op1 || !op2) return nullptr;  // 防止空指针
        int val1 = op1->getVal();
        int val2 = op2->getVal();
        switch (op) {
            case IRInstOperator::IRINST_OP_ADD_I: return new ConstInt(val1 + val2);
            case IRInstOperator::IRINST_OP_SUB_I: return new ConstInt(val1 - val2);
            case IRInstOperator::IRINST_OP_MUL_I: return new ConstInt(val1 * val2);
            case IRInstOperator::IRINST_OP_DIV_I: return new ConstInt(val1 / val2);
            case IRInstOperator::IRINST_OP_MOD_I: return new ConstInt(val1 % val2);
            // ... 其他运算
            default: return nullptr;
    }
    // 可扩展：处理浮点数、位运算等
    return nullptr;
}

// IR_optimizer.cpp
void IR_optimizer::deadCodeElimination() {
    // 1. 消除未使用的全局变量
    eliminateUnusedGlobalVariables();
    // 2. 清除未使用的函数
    eliminateUnusedFunctions();
    // 3. 消除函数中的死代码
    for (auto *func : module->getFunctionList()) {
        if (!func->isBuiltin()) {
            // 消除未使用的局部变量
            eliminateUnusedLocalVariables(func);
        }
    }
}

void IR_optimizer::eliminateUnusedFunctions() {
    auto &functions = module->getFunctionList();
    std::vector<Function*> toRemove;
    for (auto *func : functions) {
        if (!func->getIsCalled() && !func->getIsmain()) {
            toRemove.push_back(func);
        }
    }
    for (auto *func : toRemove) {
        printf("删除未使用函数: %s\n", func->getName().c_str());
        func->Delete();
        functions.erase(
            std::remove(functions.begin(), functions.end(), func),
            functions.end()
        );
    }
}

void IR_optimizer::eliminateUnusedLocalVariables(Function *func) {
    auto &Locals = func->getVarValues();
    std::vector<LocalVariable*> toRemove;
    // 收集未使用的局部变量
    for (auto *localVar : Locals) {
        if (localVar->getUseList().empty()) {
            toRemove.push_back(localVar);
        }
    }
    
    // 移除未使用的局部变量
    for (auto *localVar : toRemove) {
        Locals.erase(std::remove(Locals.begin(), Locals.end(), localVar), Locals.end());
        delete localVar; // 释放内存
    }
}

void IR_optimizer::eliminateUnusedGlobalVariables() {
    auto &globals = module->getGlobalVariables();
    std::vector<GlobalVariable*> toRemove;
    // 收集未使用的全局变量
    for (auto *globalVar : globals) {
        if (globalVar->getUseList().empty() && !globalVar->isInBSSSection()) {
            toRemove.push_back(globalVar);
        }
    }
    
    // 移除未使用的全局变量
    for (auto *globalVar : toRemove) {
        globals.erase(std::remove(globals.begin(), globals.end(), globalVar), globals.end());
        delete globalVar; // 释放内存
    }
}

void IR_optimizer::eliminateDeadCodeInFunction(Function *func) {
    auto &interCode = func->getInterCode();
    std::vector<Instruction*> toRemove;
    
    // 遍历函数中的所有指令
    for (auto *inst : interCode.getInsts()) {
        // 判断指令是否可以删除
        if (shouldRemoveInstruction(inst)) {
            toRemove.push_back(inst);
        }
    }
    
    // 移除死代码
    auto &insts = interCode.getInsts();
    insts.erase(
        std::remove_if(insts.begin(), insts.end(),
            [&](Instruction *inst) {
                if (std::find(toRemove.begin(), toRemove.end(), inst) != toRemove.end()) {
                    inst->clearOperands();
                    delete inst;
                    return true;
                }
                return false;
            }),
        insts.end());
}

bool IR_optimizer::shouldRemoveInstruction(Instruction *inst) {
    // 1. 指令没有副作用
    if (inst->hasSideEffects()) {
        return false;
    }
    
    // 2. 指令没有被其他指令使用
    if (!inst->getUseList().empty() && inst->getOp() != IRInstOperator::IRINST_OP_ASSIGN) {
        return false;
    }
    
    // 3. 不是函数出口标签等重要指令
    if (inst->isTerminator() || inst->isLabel() || inst->isEntry()) {
        return false;
    }
    
    // 4. 如果是赋值语句，检查目标是否有 use
    if (inst->getOp() == IRInstOperator::IRINST_OP_ASSIGN) {
        Value *target = inst->getOperand(0);
        bool isDeadStore = std::all_of(
            target->getUseList().begin(),
            target->getUseList().end(),
            [inst](Use *use) { return use->getUser() == inst; }
        );
        if (isDeadStore) {
            return true;
        }
    }
    
    return false;
}
// 其他优化函数暂时留空
void IR_optimizer::optimizeGlobalVariables() {}
void IR_optimizer::optimizeFunctions() {}
void IR_optimizer::functionInlining() {}
void IR_optimizer::registerAllocationOptimization() {}
